import argparse
import os
import cv2
import numpy as np
from modules.video_processor import VideoProcessor
from modules.document_detector import DocumentDetector
from modules.ocr_processor import OCRProcessor
from modules.pdf_generator import PDFGenerator
from modules.deduplicator import ImageDeduplicator
from modules.utils import evaluate_image_quality
from concurrent.futures import ThreadPoolExecutor
from modules.image_processor import ImageProcessor

def main():
    parser = argparse.ArgumentParser(description='从视频中提取帧并进行文档识别和OCR')
    
    parser.add_argument('--rotate', type=int, choices=[0, 90, 180, 270], default=0,
                      help='旋转角度，可选值：0, 90, 180, 270')
    parser.add_argument('--video', type=str, required=True,
                      help='输入视频文件的路径')
    parser.add_argument('--output', type=str, required=True,
                      help='输出文件夹路径')
    parser.add_argument('--similarity', type=float, default=0.98,
                      help='帧相似度阈值，范围0-1，默认0.98')
    parser.add_argument('--quality', type=float, default=60,
                      help='图像质量阈值，范围0-100，默认60')
    # 添加OCR去重控制参数
    parser.add_argument('--ocr-dedup', action='store_true',
                      help='是否使用OCR文本进行额外的去重，默认关闭')
    
    args = parser.parse_args()
    
    try:
        # 初始化各个模块，增加批处理大小配置
        video_proc = VideoProcessor(args.video, args.output, args.rotate, batch_size=50)
        doc_detector = DocumentDetector(batch_size=8)
        ocr_proc = OCRProcessor(batch_size=4)
        deduplicator = ImageDeduplicator(args.similarity)
        
        print("\n开始处理视频...")
        frames_with_doc = []
        total_processed = 0
        
        # 使用批量处理文档检测
        for batch_frames in video_proc.read_frames():
            if not batch_frames:
                continue
                
            # 批量检测文档
            frames_list = [frame for _, frame in batch_frames]
            if not frames_list:
                continue
                
            print(f"\n处理批次，包含 {len(frames_list)} 帧")
            doc_boxes = doc_detector.detect_batch(frames_list)
            
            # 并行处理质量评估
            with ThreadPoolExecutor(max_workers=8) as executor:
                quality_futures = [
                    executor.submit(evaluate_image_quality, frame)
                    for _, frame in batch_frames
                ]
                qualities = [future.result()[0] for future in quality_futures]
            
            # 收集符合条件的帧
            batch_valid_frames = 0
            for (idx, frame), doc_box, quality in zip(batch_frames, doc_boxes, qualities):
                # 修改判断条件，doc_box可能是numpy数组
                is_valid_box = isinstance(doc_box, np.ndarray) and doc_box.size > 0
                if is_valid_box:  # 移除质量判断，先看看检测结果
                    print(f"\n检测到文档 - 帧 {idx}")
                    print(f"文档框大小: {doc_box.shape}")
                    print(f"质量分数: {quality:.2f}")
                    
                    frames_with_doc.append({
                        'frame': frame,
                        'box': doc_box,
                        'quality': quality,
                        'index': idx
                    })
                    batch_valid_frames += 1
            
            total_processed += len(batch_frames)
            if batch_valid_frames > 0:
                print(f"\n本批次找到 {batch_valid_frames} 个有效文档")
            print(f"\r已处理 {total_processed} 帧，累计找到 {len(frames_with_doc)} 个有效文档", end="")
        
        print(f"\n\n文档检测完成，共找到 {len(frames_with_doc)} 个有效文档")
        
        if not frames_with_doc:
            print("未找到任何有效文档，程序退出")
            return
            
        # 先进行裁剪操作
        print("\n开始裁剪文档区域...")
        cropped_frames = []
        with ThreadPoolExecutor(max_workers=8) as executor:
            futures = []
            for frame_info in frames_with_doc:
                frame = frame_info['frame']
                box = frame_info['box']
                idx = frame_info['index']
                
                # 裁剪操作
                future = executor.submit(
                    save_frame,
                    frame,
                    box,
                    idx,
                    args.output
                )
                futures.append((future, frame_info['quality'], idx))
            
            # 收集裁剪结果
            for future, quality, idx in futures:
                result = future.result()
                if result:
                    cropped_frames.append({
                        'frame_path': result,
                        'quality': quality,
                        'index': idx
                    })
        
        if not cropped_frames:
            print("没有成功裁剪的图片，程序退出")
            return
            
        print(f"成功裁剪 {len(cropped_frames)} 张图片")
        
        # 对裁剪后的图片进行去重
        unique_frames = deduplicator.deduplicate_frames(cropped_frames)
        
        if not unique_frames:
            print("去重后没有剩余帧，程序退出")
            return
            
        # 获取去重后的图片路径
        output_images = [frame_info['frame_path'] for frame_info in unique_frames]
        print(f"去重完成，保留 {len(output_images)} 张图片")
        
        # OCR处理（如果启用）
        if args.ocr_dedup and output_images:
            image_info = ocr_proc.process_batch(output_images)
            final_images = ocr_proc.deduplicate(image_info)
            print(f"\n文本相似度去重完成，最终保留图片数: {len(final_images)}")
            images_to_pdf = final_images
        else:
            images_to_pdf = output_images
        
        # 生成PDF
        if images_to_pdf:
            pdf_path = os.path.join(args.output, 'combined.pdf')
            PDFGenerator.generate(images_to_pdf, pdf_path)
        
    except Exception as e:
        print(f"\n处理视频时出错: {str(e)}")
        raise  # 添加异常堆栈跟踪

def save_frame(frame, box, idx, output_folder):
    """保存单帧图片"""
    try:
        # 确保输出目录存在
        os.makedirs(output_folder, exist_ok=True)
        
        # 裁剪文档区域
        x1, y1 = np.min(box, axis=0)
        x2, y2 = np.max(box, axis=0)
        cropped = frame[int(y1):int(y2), int(x1):int(x2)]
        
        # 创建图像处理器实例
        # processor = ImageProcessor(target_size=(1654, 2339))
        processor = ImageProcessor(target_size=(660, 935))
        final=processor.process_image(cropped)
        
        # 创建帧保存目录
        frames_dir = os.path.join(output_folder, 'frames')
        os.makedirs(frames_dir, exist_ok=True)
        
        frame_name = f'frame_{idx:04d}'
        processed_path = os.path.join(frames_dir, f'{frame_name}_processed.jpg')
        cv2.imwrite(processed_path, final)
        
        print(f"\n保存处理后的帧到: {processed_path}")
        return processed_path
        
    except Exception as e:
        print(f"\n保存帧 {idx} 失败: {str(e)}")
        return None

if __name__ == "__main__":
    main()


